import cv2
import numpy as np


def RoughDeal(imagePath):
    try:
        img = cv2.imread(imagePath)
    except:
        print('can not open the image')

    # HSV空间转换并返回所需颜色的区域内容--mask

    def hsv(img, lower_hsv, upper_hsv):
        try:
            img2 = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
            mask = cv2.inRange(img2, lower_hsv, upper_hsv)
            return mask
        except:
            print('can not transform image to gray,maybe it is none')


    # 红色区域阈值

    lower_hsv = np.array([156, 43, 46])
    upper_hsv = np.array([180, 255, 255])
    mask1 = hsv(img, lower_hsv, upper_hsv)
    lower_hsv = np.array([0, 43, 46])
    upper_hsv = np.array([10, 255, 255])
    mask2 = hsv(img, lower_hsv, upper_hsv)

    try:
        mask3 = mask1 + mask2
        # cv2.imshow('mask1', mask1)
        # cv2.imshow('mask2', mask2)
        # cv2.imshow('mask3', mask3)
        # cv2.imwrite('mask3.png', mask3)
        # cv2.waitKey(0)

        #  获取清晰的和红色部分

        img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        _, thresh = cv2.threshold(img_gray, img.mean(), 255, cv2.THRESH_BINARY_INV)
        # cv2.imwrite('tep.png', img_gray * mask3)
        return img_gray * mask3
    except:
        print('the mask is none')


# 对数转换

def LinearTran(path, a=1, b=0):
    im = RoughDeal(path)
    imarray = np.array(im)

    try:
        height, width = imarray.shape
        for i in range(height):
            for j in range(width):
                aft = int(a * imarray[i, j] + b)
                if aft <= 255 and aft >= 0:
                    imarray[i, j] = aft
                elif aft > 255:
                    imarray[i, j] = 255
                else:
                    imarray[i, j] = 0
        return imarray
    except:
        print('imarray have no  attribute name shape')


# 从带背景字体图像中提取公章的函数

def ExtraceOfficialSeal(srcPath, savePath, aplha=2, belta=-20):
    grayImage = LinearTran(srcPath, a=aplha, b=belta)

    # 灰色图二值图转BGR图像
    try:
        tepimage = np.zeros(shape=grayImage.shape)
        officialseal = np.dstack((tepimage, tepimage, grayImage))
        cv2.imwrite(savePath, officialseal)
        return officialseal
    except:
        print('can not extract official seal')


if __name__ == '__main__':
    import os

    srclist = os.listdir('image/')
    for i in range(len(srclist)):
        srcpath = os.path.join('image', srclist[i])
        savepath = os.path.join('save', srclist[i])
        dirpath = os.path.join('dir', srclist[i])

        image = ExtraceOfficialSeal('D:\\temp\gongzhang\image\q.jpg', 'D:\\temp\gongzhang\save\q.jpg')
        print(srclist[i])

        im = cv2.imread(srcpath)
        try:
            img = np.concatenate((im, image), axis=1)
            cv2.imwrite('D:\\temp\gongzhang\dir\q.jpg', img)
        except:
            print('image is none')
